# -*- coding: utf-8 -*-

#    OriginDB
#    Copyright (C) 2010  Jason Power <power.jg@gmail.com>
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

from PyQt4.QtGui import *
from PyQt4 import QtCore

import openbabel
import pybel

from atom import Atom
from bond import Bond
#from minimize import Minimize

class Compound(QGraphicsItem):

	scale = 50.0
#	minObj = Minimize()
	ff = openbabel.OBForceField.FindForceField("mmff94")
	builder = openbabel.OBOp.FindType("Gen3D")


	def __init__(self, mol=None, showHydrogens=False, parent=None):
		QGraphicsItem.__init__(self, parent)

		if mol and type(mol)==pybel.Molecule:
			mol = mol.OBMol
		elif mol and type(mol)==str:
			pymol = pybel.readstring("sdf", mol)
			mol = pymol.OBMol
		if mol:
			self.mol = mol
		else:
			self.mol = openbabel.OBMol()

		self.mol.Center()

		self.atoms = []
		self.bonds = []

		self.showHydrogens = showHydrogens
		if showHydrogens:
			self.mol.AddHydrogens()
		elif not showHydrogens:
			self.mol.DeleteHydrogens()
		
		self.selected = False

		self.setFlag(QGraphicsItem.ItemIsMovable, True)
		self.setFlag(QGraphicsItem.ItemIsSelectable, True)
		try:
			self.setFlag(QGraphicsItem.ItemSendsGeometryChanges, True)
		except AttributeError:
			pass #Need this in Qt4.6 not in Qt4.5

	def addAtom(self, atom, pos):
		self.mol.BeginModify()
		newAtom = self.mol.NewAtom()
		newAtom.Duplicate(atom)
		newAtom.SetVector(self.translateToAtomCoords(pos))
		self.mol.EndModify()
		atomItem = Atom(newAtom, parent=self)
		self.atoms.append(atomItem)

		self.prepareGeometryChange()
		return newAtom

	def removeAtom(self, atom):
		print "Trying to remove", atom.GetIdx()
		idx = atom.GetIdx()
		for i in range(len(self.atoms)):
			if self.atoms[i].atom.GetIdx() == idx:
				self.scene().removeItem(self.atoms[i])
				del self.atoms[i]
				break

		pos = self.translateFromAtomCoords([atom.GetX(), atom.GetY()])

		self.mol.BeginModify()

		foundBond = True
		while foundBond:
			foundBond = False
			for i in range(len(self.bonds)):
				if self.bonds[i].bond.GetBeginAtomIdx() == atom.GetIdx() or self.bonds[i].bond.GetEndAtomIdx() == atom.GetIdx():
					self.scene().removeItem(self.bonds[i])
					print "Removing bond:", self.bonds[i].bond.GetIdx()
					self.mol.DeleteBond(self.bonds[i].bond)
					del self.bonds[i]
					foundBond = True
					break
		
		self.mol.DeleteAtom(atom)
		self.mol.EndModify()

		self.prepareGeometryChange()
		return pos

	def addBond(self, beginAtom, endAtom, typ):
		self.mol.BeginModify()
		bond = self.mol.NewBond()
		bond.SetBegin(beginAtom)
		bond.SetEnd(endAtom)
		if typ=="single":
			bond.SetBondOrder(1)
		elif typ=="double":
			bond.SetBondOrder(2)
		elif typ=="triple":
			bond.SetBondOrder(3)
		elif typ=="wedge":
			bond.SetBondOrder(1)
			bond.SetWedge()
		elif typ=="hash":
			bond.SetBondOrder(1)
			bond.SetHash()
		elif typ=="squiggle":
			bond.SetBondOrder(1)
		elif typ=="cistrans":
			bond.SetBondOrder(2)
		self.mol.EndModify()

		b = Bond(bond, parent=self)
		self.bonds.append(b)

		self.prepareGeometryChange()
		return bond

	def addMolecule(self, mol):
		for i in self.atoms:
			self.scene().removeItem(i)
		for i in self.bonds:
			self.scene().removeItem(i)
		self.atoms = []
		self.bonds = []

		self.mol += mol
		
		self.mol.SetDimension(2)
		if not self.builder.Do(self.mol):
			raise Exception("Builder failed")
		self.mol.DeleteHydrogens()
		
		self.addToScene()
		return [], []

	def removeBond(self, bond):
		print "Trying to remove", bond.GetIdx()
		idx = bond.GetIdx()
		for i in range(len(self.bonds)):
			if self.bonds[i].bond.GetIdx() == idx:
				self.scene().removeItem(self.bonds[i])
				del self.bonds[i]
				break
		
		beginAtom = bond.GetBeginAtom()
		endAtom = bond.GetEndAtom()

		order = bond.GetBO()
		typ = ""
		if order==1:
			if bond.IsHash():
				typ = "hash"
			elif bond.IsWedge():
				typ = "wedge"
			else:
				typ = "single"
		elif order==2:
			typ = "double"
		elif order==3:
			typ = "triple"

		self.mol.BeginModify()
		self.mol.DeleteBond(bond)
		self.mol.EndModify()

		self.prepareGeometryChange()
		print "New len of bonds:", len(self.bonds)
		return beginAtom, endAtom, typ

	def setShowHydrogens(self, show):
		self.showHyrdogens = show
		if show and not self.mol.HasHydrogensAdded():
			self.mol.AddHydrogens()
		elif not show and self.mol.HasHydrogensAdded():
			self.mol.DeleteHydrogens()

		for i in self.atoms:
			self.scene().removeItem(i)
		for i in self.bonds:
			self.scene().removeItem(i)
		self.atoms = []
		self.bonds = []
		self.addToScene()
		self.prepareGeometryChange()

	def translateFromAtomCoords(self, coords):
		""" @coords is a 2 or 3 tuple
			Returns a QPointF
		"""
		return QtCore.QPointF(coords[0]*self.scale, coords[1]*self.scale)

	def translateToAtomCoords(self, point):
		""" @point is a QPointF
			Returns a openbabel vector3
		"""
		return pybel.ob.vector3(point.x()/self.scale, point.y()/self.scale, 0)

#	def minimize(self):
		
#		self.minObj.initialize(self.mol)
		
#		if self.minObj.optimizeCoords():
#			print "moving"
#		else:
#			print self.minObj.endReason
#		self.minObj.finalize()

#		for atom in openbabel.OBMolAtomIter(self.mol):
#			print "Atom", atom.GetIdx(), atom.GetAtomicNum(), ":", atom.GetX(), atom.GetY(), atom.GetZ()

#		for atom in self.atoms:
#			atom.updatePos()

		

	def getAtom(self, idx):
		if idx > self.mol.NumAtoms():
			raise ValueError("Idx is too large, bug!")
		return self.mol.GetAtom(idx)

	def addToScene(self):
		for atom in openbabel.OBMolAtomIter(self.mol):
			a = Atom(atom, parent=self)
			self.atoms.append(a)

		for bond in openbabel.OBMolBondIter(self.mol):
			b = Bond(bond, parent=self)
			self.bonds.append(b)

	def updateBonds(self, atom):
		for bond in self.bonds:
			if bond.bond.GetBeginAtomIdx()==atom.GetIdx() or bond.bond.GetEndAtomIdx()==atom.GetIdx():
				bond.atomMoved()
				
		self.prepareGeometryChange()

	def itemChange(self, change, value):
		if change == QGraphicsItem.ItemSceneHasChanged and value:
			self.addToScene()
		return QGraphicsItem.itemChange(self, change, value)

	def setSelected(self, selected):
		self.selected = selected
		self.update()

	def boundingRect(self):
		return self.childrenBoundingRect().adjusted(-2,-2,2,2)

	def paint(self, painter, option, widget):
		if self.selected:
			painter.drawRect(self.childrenBoundingRect())
